import mplfinance as mpf
import numpy as np
import pandas as pd


# 很多字体不支持
# 标题格式，字体为中文字体，颜色为黑色，粗体，水平中心对齐
title_font = {
  # 'fontname': 'Times New Roman, SimSun', 
  'size':     '16',
  'color':    'black',
  'weight':   'bold',
  'va':       'bottom',
  'ha':       'center'
}
# 红色数字格式（显示开盘收盘价）粗体红色24号字
large_red_font = {
  # 'fontname': 'Times New Roman, SimSun',
  'size':     '24',
  'color':    'red',
  'weight':   'bold',
  'va':       'bottom'
}
# 绿色数字格式（显示开盘收盘价）粗体绿色24号字
large_green_font = {
  # 'fontname': 'Times New Roman, SimSun',
  'size':     '24',
  'color':    'green',
  'weight':   'bold',
  'va':       'bottom'
}
# 小数字格式（显示其他价格信息）粗体红色12号字
small_red_font = {
  # 'fontname': 'Times New Roman, SimSun',
  'size':     '12',
  'color':    'red',
  'weight':   'bold',
  'va':       'bottom'
}
# 小数字格式（显示其他价格信息）粗体绿色12号字
small_green_font = {
  # 'fontname': 'Times New Roman, SimSun',
  'size':     '12',
  'color':    'green',
  'weight':   'bold',
  'va':       'bottom'
}
# 标签格式，可以显示中文，普通黑色12号字
normal_label_font = {
  # 'fontname': 'Microsoft PhagsPa',
  'size':     '12',
  'color':    'black',
  'va':       'bottom',
  'ha':       'right'
}
# 普通文本格式，普通黑色12号字
normal_font = {
  # 'fontname': 'Times New Roman, SimSun',
  'size':     '12',
  'color':    'black',
  'va':       'bottom',
  'ha':       'left'
}

# 定义一个交互K线图类
class InterCandle: 
  def __init__(self, data, my_style):
    # 初始化交互式K线图对象，历史数据作为唯一的参数用于初始化对象
    self.data = data
    self.style = my_style
    # 设置初始化的K线图显示区间起点为0，即显示第0到第99个交易日的数据（前100个数据）
    self.idx_start = 0
    
    # 初始化figure对象，在figure上建立三个Axes对象并分别设置好它们的位置和基本属性
    self.fig = mpf.figure(style=my_style, figsize=(12, 8), facecolor=(0.82, 0.83, 0.85))
    fig = self.fig
    self.ax1 = fig.add_axes([0.08, 0.25, 0.88, 0.60])
    self.ax2 = fig.add_axes([0.08, 0.15, 0.88, 0.10], sharex=self.ax1)
    self.ax2.set_ylabel('volume')
    self.ax3 = fig.add_axes([0.08, 0.05, 0.88, 0.10], sharex=self.ax1)
    self.ax3.set_ylabel('macd')
    # 初始化figure对象，在figure上预先放置文本并设置格式，文本内容根据需要显示的数据实时更新
    # 初始化时，所有的价格数据都显示为空字符串
    last_data = data.iloc[-1]

    self.t1 = fig.text(0.50, 0.94, 'sh.600508:--上海能源', **title_font)
    self.t2 = fig.text(0.12, 0.90, '开/收: ', **normal_label_font)
    self.t3 = fig.text(0.13, 0.90, f'{np.round(last_data["open"], 3)} / {np.round(last_data["close"], 3)}', **large_red_font)
    self.t4 = fig.text(0.12, 0.86, '最后交易日期:', **normal_label_font)
    self.t5 = fig.text(0.2, 0.86, f'{last_data.name.date()}', **normal_label_font)
    self.t6 = fig.text(0.40, 0.90, '高: ', **normal_label_font)
    self.t7 = fig.text(0.42, 0.90, f'{last_data["high"]}', **small_red_font)
    self.t8 = fig.text(0.40, 0.86, '低: ', **normal_label_font) 
    self.t9 = fig.text(0.42, 0.86, f'{last_data["low"]}', **small_green_font)
    self.t10 = fig.text(0.55, 0.90, '量(股): ', **normal_label_font)
    self.t11 = fig.text(0.57, 0.90, f'{last_data["volume"]}', **normal_font)
    self.t12 = fig.text(0.55, 0.86, '额(元): ', **normal_label_font)
    self.t13 = fig.text(0.57, 0.86, f'{last_data["amount"]}', **normal_font)
    self.t14 = fig.text(0.70, 0.90, '涨停: ', **normal_label_font)
    self.t15 = fig.text(0.70, 0.86, '跌停: ', **normal_label_font)
    self.t16 = fig.text(0.85, 0.90, '均价: ', **normal_label_font)
    self.t17 = fig.text(0.85, 0.86, '昨收: ', **normal_label_font)
    self.t18 = fig.text(0.89, 0.86, f'{last_data["preclose"]}', **normal_font)


    # 鼠标按键状态，False为按键未按下，True为按键按下
    self.pressed = False
    # 鼠标按下时的x坐标
    self.xpress = None

    # 以下部分代码与本节开头的完全相同，为节省篇幅，略去不表
		# 下面的代码在__init__()中，告诉matplotlib哪些回调函数用于响应哪些事件
		# 鼠标按下事件与self.on_press回调函数绑定
    fig.canvas.mpl_connect('button_press_event', self.on_press)
    # 鼠标按键释放事件与self.on_release回调函数绑定
    fig.canvas.mpl_connect('button_release_event', self.on_release)
    # 鼠标移动事件与self.on_motion回调函数绑定
    fig.canvas.mpl_connect('motion_notify_event', self.on_motion)


    # 其他代码未发生变化，因而省略
    self.idx_range = 100  # 控制K线图的显示范围大小
    # 将新增的回调函数on_scroll与鼠标滚轮事件绑定起来
    self.fig.canvas.mpl_connect('scroll_event', self.on_scroll)

  def on_press(self, event):
    # 当鼠标按键按下时，调用该函数，event为事件信息，是一个dict对象，包含事件相关的信息
    # 如坐标、按键类型、是否在某个Axes对象内等等
    # event.inaxes可用于判断事件发生时，鼠标是否在某个Axes内，在这里我们指定，只有鼠
    # 标在ax1内时，才能平移K线图，否则就退出事件处理函数
    if not event.inaxes == self.ax1:
      return
    # 检查是否按下了鼠标左键，如果不是左键，同样退出事件处理函数
    if event.button != 1:
      return
    # 如果鼠标在ax1范围内，且按下了左键，条件满足，设置鼠标状态为pressed
    self.pressed = True
    # 同时记录鼠标按下时的x坐标，退出函数，等待鼠标移动事件发生
    self.xpress = event.xdata
	
  # 鼠标移动事件处理
  def on_motion(self, event):
    # 如果鼠标按键没有按下pressed == False，则什么都不做，退出处理函数
    if not self.pressed:
      return
    # 如果移动出了ax1的范围，也退出处理函数
    if not event.inaxes == self.ax1:
      return
    # 如果鼠标在ax1范围内，且左键按下，则开始计算dx，并根据dx计算新的K线图起点
    dx = int(event.xdata - self.xpress)
    # 前面介绍过了，新的起点N(new) = N - dx
    new_start = self.idx_start - dx
    # 设定平移的左右界限，如果平移后超出界限，则不再平移
    if new_start <= 0:
      new_start = 0
    if new_start >= len(self.data) - 100:
      new_start = len(self.data) - 100
    # 清除各个图表Axes中的内容，准备以新的起点重新绘制
    self.ax1.clear()
    self.ax2.clear()
    self.ax3.clear()
    # 更新图表上的文字、以新的起点开始绘制K线图
    self.refresh_texts(self.data.iloc[new_start])
    self.refresh_plot(new_start, self.idx_range)
	
	# 鼠标按键释放
  def on_release(self, event):
  # 按键释放后，设置鼠标的pressed为False
    self.pressed = False
    # 此时别忘了最后一次更新K线图的起点，否则下次拖拽的时候就不会从这次的起点开始移动了
    dx = int(event.xdata - self.xpress)
    self.idx_start -= dx
    if self.idx_start <= 0:
      self.idx_start = 0
    if self.idx_start >= len(self.data) - 100:
      self.idx_start = len(self.data) - 100

  def on_scroll(self, event):
    # 仅当鼠标滚轮在axes1范围内滚动时起作用
    if event.inaxes != self.ax1:
      return
    if event.button == 'down':
      # 缩小20%显示范围
      scale_factor = 0.8
    if event.button == 'up':
      # 放大20%显示范围
      scale_factor = 1.2
		# 设置K线的显示范围大小
    self.idx_range = int(self.idx_range * scale_factor)
    # 限定可以显示的K线图的范围，最少不能少于30个交易日，最大不能超过当前位置与
    # K线数据总长度的差
    data_length = len(self.data)
    if self.idx_range >= data_length - self.idx_start:
      self.idx_range = data_length - self.idx_start
    if self.idx_range <= 30:
      self.idx_range = 30 
		# 更新图表（注意因为多了一个参数idx_range，refresh_plot函数也有所改动）
    self.ax1.clear()
    self.ax2.clear()
    self.ax3.clear()
    self.refresh_texts(self.data.iloc[self.idx_start])
    self.refresh_plot(self.idx_start, self.idx_range)


  def refresh_plot(self, idx_start, idx_range):
    """ 根据最新的参数，重新绘制整个图表
    """
    all_data = self.data
    plot_data = all_data.iloc[idx_start: idx_start + idx_range]

    ap = []
    # 添加K线图重叠均线
    ap.append(mpf.make_addplot(plot_data[['MA5', 'MA10', 'MA20', 'MA60']], ax=self.ax1))
    # 添加指标MACD
    ap.append(mpf.make_addplot(plot_data[['macd-m', 'macd-s']], ax=self.ax3))
    bar_r = np.where(plot_data['macd-h'] > 0, plot_data['macd-h'], 0)
    bar_g = np.where(plot_data['macd-h'] <= 0, plot_data['macd-h'], 0)
    ap.append(mpf.make_addplot(bar_r, type='bar', color='red', ax=self.ax3))
    ap.append(mpf.make_addplot(bar_g, type='bar', color='green', ax=self.ax3))
    # 绘制图表
    mpf.plot(
      plot_data,
      ax=self.ax1,
      volume=self.ax2,
      addplot=ap,
      type='candle',
      style=self.style,
      datetime_format='%Y-%m',
      xrotation=0
    )
    self.fig.show()


# 参数名称	参数描述	算法说明
# date	交易所行情日期	
# code	证券代码	
# open	开盘价	
# high	最高价	
# low	最低价	
# close	收盘价	
# preclose	前收盘价	见表格下方详细说明
# volume	成交量（累计 单位：股）	
# amount	成交额（单位：人民币元）	
# adjustflag	复权状态(1：后复权， 2：前复权，3：不复权）	
# turn	换手率	[指定交易日的成交量(股)/指定交易日的股票的流通股总股数(股)]*100%
# tradestatus	交易状态(1：正常交易 0：停牌）	
# pctChg	涨跌幅（百分比）	日涨跌幅=[(指定交易日的收盘价-指定交易日前收盘价)/指定交易日前收盘价]*100%
# peTTM	滚动市盈率	(指定交易日的股票收盘价/指定交易日的每股盈余TTM)=(指定交易日的股票收盘价*截至当日公司总股本)/归属母公司股东净利润TTM
# pbMRQ	市净率	(指定交易日的股票收盘价/指定交易日的每股净资产)=总市值/(最近披露的归属母公司股东的权益-其他权益工具)
# psTTM	滚动市销率	(指定交易日的股票收盘价/指定交易日的每股销售额)=(指定交易日的股票收盘价*截至当日公司总股本)/营业总收入TTM
# pcfNcfTTM	滚动市现率	(指定交易日的股票收盘价/指定交易日的每股现金流TTM)=(指定交易日的股票收盘价*截至当日公司总股本)/现金以及现金等价物净增加额TTM
# isST	是否ST股，1是，0否	

  def refresh_texts(self, display_data):
    """ 更新K线图上的价格文本
    """
    # display_data是一个交易日内的所有数据，将这些数据分别填入figure对象上的文本中
    # print(display_data)
    self.t3.set_text(f'{np.round(display_data["open"], 3)} / {np.round(display_data["close"], 3)}')
    self.t5.set_text(f'{display_data.name.date()}')
    self.t7.set_text( f'{display_data["high"]}')
    self.t9.set_text(f'{display_data["low"]}')
    self.t11.set_text(f'{display_data["volume"]}')
    self.t13.set_text(f'{display_data["amount"]}')
    self.t18.set_text(f'{display_data["preclose"]}')

    # 根据本交易日的价格变动值确定开盘价、收盘价的显示颜色
    if display_data['pctChg'] > 0:  # 如果今日变动额大于0，即今天价格高于昨天，今天价格显示为红色
        close_number_color = 'red'
    elif display_data['pctChg'] < 0:  # 如果今日变动额小于0，即今天价格低于昨天，今天价格显示为绿色
        close_number_color = 'green'
    else:
        close_number_color = 'black'
    self.t1.set_color(close_number_color)
    self.t2.set_color(close_number_color)
    self.t3.set_color(close_number_color)


# 获取数据，处理数据
def read_data(filename):
  df = pd.read_csv('./files/history_A_stock_k_data.csv')
  df['Open'] = df['open']
  df['Close'] = df['close']
  df['High'] = df['high']
  df['Low'] = df['low']
  df['Volume'] = df['volume']
  df['Date'] = df['date']

  df['Date'] = pd.to_datetime(df['Date'])
  df.set_index(['Date'], inplace=True)
  return df

# 获取数据
df = read_data('./files/test02.csv')

df['MA5'] = df['Close'].rolling(window=5).mean()
df['MA10'] = df['Close'].rolling(window=10).mean()
df['MA20'] = df['Close'].rolling(window=20).mean()
df['MA60'] = df['Close'].rolling(window=60).mean()


df['Short_MA'] = df['Close'].rolling(window=12).mean()
df['Long_MA'] = df['Close'].rolling(window=26).mean()
df['DIF'] = df['Short_MA'] - df['Long_MA']        # 快线
df['DEA'] = df['DIF'].rolling(window=9).mean()    # 慢线
df['MACD'] = df['DIF'] - df['DEA']                # 柱子

# df['macd-s'] = df['DEA']
# df['macd-m'] = df['DIF']

df['macd-s'] = df['DIF']
df['macd-m'] = df['DEA']
df['macd-h'] = df['MACD']

# 设置线元素的颜色
my_color = mpf.make_marketcolors(
  up='r',
  down='g',
  edge='inherit',
  wick='inherit',
  volume='inherit'
)

# 自定义风格
my_style = mpf.make_mpf_style(
  marketcolors=my_color,
  figcolor='(0.82, 0.83, 0.85)',
  gridcolor='(0.82, 0.83, 0.85)',
  rc={'font.family': 'SimHei', 'axes.unicode_minus': 'False'}
)


# 创建一个InterCandle对象，显示风格为前面定义好的my_style风格（即中国股市惯例风格）
candle = InterCandle(df, my_style)
# 更新图表上的文本，显示第100个交易日的K线数据
candle.refresh_texts(df.iloc[100])
# 更新显示第100天开始的K线图
candle.refresh_plot(100, 100)
input("Press Enter to close...")

